import cv2
import numpy as np
import networkx as nx
import random
import matplotlib.pyplot as plt
from enum import Enum
from typing import Tuple, Optional, List, Dict
from collections import defaultdict
import os


class Direction(Enum):
    LEFT = 0
    UP = 1
    RIGHT = 2
    DOWN = 3


# 方向到颜色的映射（RGBA格式）
DIRECTION_COLORS = {
    Direction.LEFT: (0, 0, 255, 100),  # 蓝色（半透明）
    Direction.UP: (0, 255, 0, 100),  # 绿色
    Direction.RIGHT: (255, 0, 0, 100),  # 红色
    Direction.DOWN: (255, 255, 0, 100)  # 黄色
}


class PixelWalker:
    def __init__(self, image_path: str):
        self.img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        if self.img is None:
            raise ValueError("Image not found")

        if self.img.shape[2] == 3:
            self.img = cv2.cvtColor(self.img, cv2.COLOR_BGR2BGRA)

        self.height, self.width = self.img.shape[:2]
        self.graph = nx.DiGraph()
        self.visited = set()
        self.reset_walker()
        self.iteration = 0
        os.makedirs("iterations", exist_ok=True)
        self.current_degree = 1  # 新增属性，记录当前升维程度，初始为1（即不升维）
        self.original_channels = ['R', 'G', 'B', 'A']  # 原始通道名称
        self.dimension_names = self.original_channels.copy()  # 当前所有维度的表达式

    def reset_walker(self):
        """随机初始化漫游器位置"""
        self.current_pos = (
            random.randint(0, self.width - 1),
            random.randint(0, self.height - 1)
        )
        self.current_pixel = self.get_pixel(*self.current_pos)
        start_key = self.pos_key(self.current_pos)
        self.graph.add_node(start_key,
                            pixel=self.current_pixel,
                            pos=self.current_pos)
        self.current_node = start_key
        self.visited.add(self.current_pos)

    def pos_key(self, pos: Tuple[int, int]) -> str:
        return f"{pos[0]}_{pos[1]}"

    def get_pixel(self, x: int, y: int) -> Tuple[int,...]:
        return tuple(self.img[y, x])

    def pixel_equal(self, p1, p2, threshold: int = 5) -> bool:
        """通用像素判等器，可以处理不同维度的像素"""
        # 确保两个像素维度相同
        if len(p1)!= len(p2):
            return False
        return all(abs(int(c1) - int(c2)) <= threshold for c1, c2 in zip(p1, p2))

    def polynomial_expansion(self, pixel: Tuple, degree: int = 2) -> Tuple:
        """多项式升维函数"""
        expanded = list(pixel)
        for d in range(2, degree + 1):
            for c in pixel:
                expanded.append(c ** d)
        return tuple(expanded)

    def expand_image_dimensions(self):
        """将整个图像进行逐维升维"""
        channel_idx = self.current_degree - 1  # 原始通道索引（0=R, 1=G, 2=B, 3=A）
        exponent = self.current_degree + 1       # 指数 = 当前度数 + 1（第一次为2，第二次为3）
        new_dim = f"{self.original_channels[channel_idx % len(self.original_channels)]}^{exponent}"

        new_img = np.zeros((self.height, self.width, self.img.shape[2] + 1), dtype=np.uint8)
        for y in range(self.height):
            for x in range(self.width):
                pixel = self.get_pixel(x, y)
                new_pixel = list(pixel)
                new_pixel.append(pixel[self.current_degree - 1] ** (self.current_degree + 1))  # 逐维扩展
                new_img[y, x] = tuple(new_pixel)
        self.img = new_img

        # 更新维度名称和度数
        self.dimension_names.append(new_dim)
        self.current_degree += 1

    def find_existing_node(self, pixel: Tuple) -> Optional[str]:
        for node, data in self.graph.nodes(data=True):
            if self.pixel_equal(data['pixel'], pixel):
                return node
        return None

    def move(self, direction: Direction) -> Tuple[int, int]:
        x, y = self.current_pos
        dx, dy = {
            Direction.LEFT: (-1, 0),
            Direction.UP: (0, 1),
            Direction.RIGHT: (1, 0),
            Direction.DOWN: (0, -1)
        }[direction]
        return (x + dx, y + dy)

    def is_out_of_bounds(self, pos: Tuple[int, int]) -> bool:
        x, y = pos
        return x < 0 or x >= self.width or y < 0 or y >= self.height

    def get_opposite_direction(self, direction: Direction) -> Direction:
        """获取相反方向"""
        return Direction((direction.value + 2) % 4)

    def walk_step(self) -> bool:
        """完全随机漫游策略（记录边的方向）"""
        # 随机选择初始方向
        initial_dir = random.choice(list(Direction))
        new_pos = self.move(initial_dir)

        # 越界处理：使用反方向并更新实际移动方向
        actual_dir = initial_dir  # 初始化实际方向为初始方向
        if self.is_out_of_bounds(new_pos):
            actual_dir = self.get_opposite_direction(initial_dir)  # 实际方向改为反方向
            new_pos = self.move(actual_dir)
            if self.is_out_of_bounds(new_pos):
                return True  # 反方向仍越界，保持原位

        # 更新状态（使用实际方向）
        self.current_pos = new_pos
        new_pixel = self.get_pixel(*new_pos)
        self.visited.add(new_pos)

        # 处理图节点（使用实际方向记录边）
        existing_node = self.find_existing_node(new_pixel)
        new_node_key = self.pos_key(new_pos)

        if existing_node and existing_node != new_node_key:
            self.graph.add_edge(self.current_node, existing_node, direction=actual_dir)  # 使用实际方向
            self.current_node = existing_node
        else:
            self.graph.add_node(new_node_key, pixel=new_pixel, pos=new_pos)
            self.graph.add_edge(self.current_node, new_node_key, direction=actual_dir)  # 使用实际方向
            self.current_node = new_node_key

        self.current_pixel = new_pixel
        return True

    def has_multiple_same_color_edges(self) -> bool:
        """检查是否存在相同颜色的出度大于1的顶点"""
        color_outdegree = defaultdict(lambda: defaultdict(int))

        for u, v, data in self.graph.edges(data=True):
            direction = data['direction']
            pixel = self.graph.nodes[u]['pixel']
            color_key = tuple(pixel)  # 使用像素值作为颜色键

            color_outdegree[color_key][direction] += 1

        # 检查是否有任何颜色在任何方向上有超过1条边
        for color, directions in color_outdegree.items():
            for dir, count in directions.items():
                if count > 1:
                    return True
        return False

    def visualize_graph(self, output_path: Optional[str] = None):
        """带方向颜色标记的嵌入可视化"""
        if output_path is None:
            output_path = f"iterations/iteration_{self.iteration}.png"

        plt.figure(figsize=(10, 10))

        # 绘制原图
        if self.img.shape[2] == 4:
            img_rgb = cv2.cvtColor(self.img, cv2.COLOR_BGRA2RGBA)
        else:
            # 对于高维像素，只显示前4个通道
            img_rgb = cv2.cvtColor(self.img[:, :, :4], cv2.COLOR_BGRA2RGBA)
        plt.imshow(img_rgb)

        # 节点位置（原始坐标）
        pos = {node: data['pos'] for node, data in self.graph.nodes(data=True)}

        # 按方向绘制边
        edge_colors = []
        for u, v, data in self.graph.edges(data=True):
            direction = data['direction']
            # 将颜色转换为0-1范围（matplotlib要求）
            r, g, b, a = DIRECTION_COLORS[direction]
            edge_colors.append((r / 255, g / 255, b / 255, a / 255))

        # 绘制有向边（带方向颜色）
        nx.draw_networkx_edges(
            self.graph, pos,
            edge_color=edge_colors,
            arrowsize=15,
            width=2,
            arrowstyle='->,head_width=0.3,head_length=0.5',
            connectionstyle='arc3,rad=0.1'
        )

        # 绘制节点（淡化显示）
        nx.draw_networkx_nodes(
            self.graph, pos,
            node_size=10,
            node_color='white',
            alpha=0.3,
            edgecolors='black'
        )

        # 标记起点和终点（突出显示，不依赖方向）
        if len(self.graph.nodes) > 0:
            nodes = list(self.graph.nodes)
            plt.scatter(
                pos[nodes[0]][0], pos[nodes[0]][1],
                s=50, c='green', marker='o', edgecolors='darkgreen',
                label='Start'
            )
            plt.scatter(
                pos[nodes[-1]][0], pos[nodes[-1]][1],
                s=50, c='red', marker='x',
                label='End'
            )

        # 坐标轴设置
        plt.xlim(-0.5, self.width - 0.5)
        plt.ylim(-0.5, self.height - 0.5)
        plt.axis('off')
        plt.gca().set_aspect('equal', adjustable='box')

        # 保存图像
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0, dpi=300)
        plt.close()
        print(f"Direction-colored graph saved to {output_path}")

    def run_iterations(self, max_steps_per_iter=1000, max_iterations=10):
        """运行迭代过程直到满足条件或达到最大迭代次数"""
        for self.iteration in range(max_iterations):
            # 重置漫游器和图
            self.graph = nx.DiGraph()
            self.visited = set()
            self.reset_walker()

            # 执行漫游
            for _ in range(max_steps_per_iter):
                self.walk_step()

            # 输出当前迭代的张成组
            print(f"Iteration {self.iteration + 1} Spanning Set: {self.dimension_names}")

            # 可视化当前迭代
            self.visualize_graph()

            # 检查条件
            if not self.has_multiple_same_color_edges():
                print(f"Converged! Final Spanning Set: {self.dimension_names}\n")
                break

            # 多项式升维
            print(f"Expanding dimension to add {self.original_channels[(self.current_degree - 1) % len(self.original_channels)]}^{self.current_degree + 1}\n")
            self.expand_image_dimensions()
        else:
            print(f"Reached max iterations. Final Spanning Set: {self.dimension_names}\n")


# 使用示例
if __name__ == "__main__":
    walker = PixelWalker("digit_image.png")
    walker.run_iterations(max_steps_per_iter=1000, max_iterations=100)